package com.health.trans.annotation;

import com.health.trans.util.EntityOperations;
import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.BeanDefinitionHolder;
import org.springframework.beans.factory.config.ConstructorArgumentValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.autoconfigure.domain.EntityScanPackages;
import org.springframework.context.ResourceLoaderAware;
import org.springframework.context.annotation.ClassPathBeanDefinitionScanner;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.annotation.AnnotationAttributes;
import org.springframework.core.io.ResourceLoader;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.core.type.filter.AnnotationTypeFilter;
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
import org.springframework.data.jpa.repository.config.JpaRepositoryConfigExtension;
import org.springframework.data.repository.config.RepositoryBeanDefinitionRegistrarSupport;
import org.springframework.data.repository.config.RepositoryConfigurationExtension;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;
import org.springframework.util.StringUtils;

import javax.persistence.Entity;
import java.lang.annotation.Annotation;
import java.util.*;

/**
 * @author nqlz
 */
@Configuration
@EnableBatchProcessing
public class BatchRegister {

    public static class MapperBeanDefinitionScanner extends ClassPathBeanDefinitionScanner {

        public MapperBeanDefinitionScanner(BeanDefinitionRegistry registry, boolean useDefaultFilters) {
            super(registry, useDefaultFilters);
        }

        protected void registerFilters() {
            addIncludeFilter(new AnnotationTypeFilter(Entity.class));
        }

        @Override
        protected Set<BeanDefinitionHolder> doScan(String... basePackages) {
            return super.doScan(basePackages);
        }
    }


    static class RegistrarJpa extends RepositoryBeanDefinitionRegistrarSupport {
        @Override
        protected Class<? extends Annotation> getAnnotation() {
            return EnableJpaRepositories.class;
        }

        @Override
        protected RepositoryConfigurationExtension getExtension() {
            return new JpaRepositoryConfigExtension();
        }
    }

    static class Registrar implements ImportBeanDefinitionRegistrar, ResourceLoaderAware {
        public ResourceLoader resourceLoader;

        @Override
        public void setResourceLoader(ResourceLoader resourceLoader) {
            this.resourceLoader = resourceLoader;
        }

        @Override
        public void registerBeanDefinitions(AnnotationMetadata metadata, BeanDefinitionRegistry registry) {
            Set<String> packagesToScan = getPackagesToScan(metadata);
            EntityOperations.packageNames = packagesToScan;
            MapperBeanDefinitionScanner scanner = new MapperBeanDefinitionScanner(registry, false);
            scanner.setResourceLoader(resourceLoader);
            scanner.registerFilters();
            scanner.addIncludeFilter(new AnnotationTypeFilter(BatchJob.class));
            scanner.doScan(packagesToScan.toArray(new String[]{}));
            register(registry, packagesToScan);
        }

        private Set<String> getPackagesToScan(AnnotationMetadata metadata) {
            AnnotationAttributes attributes = AnnotationAttributes
                    .fromMap(metadata.getAnnotationAttributes(EnableBatch.class.getName()));
            assert attributes != null;
            String[] basePackages = attributes.getStringArray("entityPackages");
            Class<?>[] basePackageClasses = attributes.getClassArray("entityPackageClasses");
            Set<String> packagesToScan = new LinkedHashSet<>(Arrays.asList(basePackages));
            for (Class<?> basePackageClass : basePackageClasses) {
                packagesToScan.add(ClassUtils.getPackageName(basePackageClass));
            }
            if (packagesToScan.isEmpty()) {
                String packageName = ClassUtils.getPackageName(metadata.getClassName());
                Assert.state(!StringUtils.isEmpty(packageName), "@EntityScan cannot be used with the default package");
                return Collections.singleton(packageName);
            }
            return packagesToScan;
        }

        public static void register(BeanDefinitionRegistry registry, Collection<String> packageNames) {
            Assert.notNull(registry, "Registry must not be null");
            Assert.notNull(packageNames, "PackageNames must not be null");
            if (registry.containsBeanDefinition(EntityScanPackages.class.getName())) {
                BeanDefinition beanDefinition = registry.getBeanDefinition(EntityScanPackages.class.getName());
                ConstructorArgumentValues constructorArguments = beanDefinition.getConstructorArgumentValues();
                constructorArguments.addIndexedArgumentValue(0, addPackageNames(constructorArguments, packageNames));
            } else {
                GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
                beanDefinition.setBeanClass(EntityScanPackages.class);
                beanDefinition.getConstructorArgumentValues().addIndexedArgumentValue(0,
                        StringUtils.toStringArray(packageNames));
                beanDefinition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
                registry.registerBeanDefinition(EntityScanPackages.class.getName(), beanDefinition);
            }
        }

        private static String[] addPackageNames(ConstructorArgumentValues constructorArguments,
                                                Collection<String> packageNames) {
            String[] existing = (String[]) constructorArguments.getIndexedArgumentValue(0, String[].class).getValue();
            Set<String> merged = new LinkedHashSet<>();
            merged.addAll(Arrays.asList(existing));
            merged.addAll(packageNames);
            return StringUtils.toStringArray(merged);
        }
    }

}
